#include "Eigen/Dense"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/shape_inference.h"

namespace tensorflow {
REGISTER_OP("PsOp")
  .Attr("op: string")
  .Input("key: string")
  .Input("value: float32")
  .Output("updated: float32")
  .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(1));
      return Status::OK();
  });

REGISTER_OP("PsMyRank")
  .Output("rank: int32")
  .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->Scalar());
      return Status::OK();
  });
}

namespace ps {
namespace tensorflow {

extern void
ps_pull_data(std::vector<uint64_t> const& keys,
             std::vector<float>& data,
             std::vector<int>& lens);

extern void
ps_push_data(std::vector<uint64_t> const& keys,
             std::vector<float> const& data,
             std::vector<int> const& lens);
extern void
ps_push_pull_data(std::vector<uint64_t> const& keys,
                  std::vector<float> const& data,
                  std::vector<float>& out,
                  std::vector<int>& lens);

extern int
ps_my_rank(void);

class PsMyRankOp : public ::tensorflow::OpKernel
{
  public:
    explicit PsMyRankOp(::tensorflow::OpKernelConstruction* context)
      : OpKernel(context)
    {}

    virtual void Compute(::tensorflow::OpKernelContext* context) override
    {
        ::tensorflow::Tensor* output;
        ::tensorflow::TensorShape shape;
        OP_REQUIRES_OK(context, context->allocate_output(0, shape, &output));
        auto rank = ps_my_rank();
        output->flat<int32_t>().setConstant(rank);
    }
};

class PsOp : public ::tensorflow::OpKernel
{
  public:
    explicit PsOp(::tensorflow::OpKernelConstruction* context)
      : OpKernel(context)
    {
        std::string op;
        OP_REQUIRES_OK(context, context->GetAttr("op", &op));
        assert(op == "Push" || op == "Pull" || op == "PushPull");

        if (op == "Push") {
            _op = Push;
        } else if (op == "Pull") {
            _op = Pull;
        } else {
            _op = PushPull;
        }
    }

    virtual void Compute(::tensorflow::OpKernelContext* context) override
    {
        ::tensorflow::Tensor const& key_tensor = context->input(0);
        ::tensorflow::Tensor const& values_tensor = context->input(1);

        assert(key_tensor.shape().dims() == 0);

        auto keys = key_tensor.flat<std::string>();
        auto values = values_tensor.flat<float>();
        ::tensorflow::hash<std::string> hash_fn;

        std::vector<uint64_t> ps_key = { hash_fn(keys(0)) };
        std::vector<float> vals;
        std::vector<int> lens = { (int)values.size() };
        std::vector<float> out;

        for (auto i = 0; i < values.size(); ++i) {
            vals.push_back(values(i));
        }

        if (_op == PushPull) {
            ps_push_pull_data(ps_key, vals, out, lens);
            assert(!lens.empty() && lens[0] == (int)values.size() &&
                   lens[0] == out.size());
        } else if (_op == Push) {
            ps_push_data(ps_key, vals, lens);
        } else {
            lens.clear();
            out.clear();
            ps_pull_data(ps_key, out, lens);
            if (lens.empty() || out.empty()) {
                context->CtxFailure(::tensorflow::Status(
                  ::tensorflow::error::DATA_LOSS, "failed or not exist"));
                return;
            }
        }

        ::tensorflow::Tensor* output_tensor;
        OP_REQUIRES_OK(
          context,
          context->allocate_output(0, values_tensor.shape(), &output_tensor));
        if (out.empty()) {
            auto ok =
              output_tensor->CopyFrom(values_tensor, values_tensor.shape());
            if (!ok) {
                context->CtxFailureWithWarning(
                  __FILE__,
                  __LINE__,
                  ::tensorflow::Status(::tensorflow::error::UNKNOWN,
                                       "Copy From input failed."));
            }
            return;
        }

        auto pout = output_tensor->flat<float>();
        for (auto i = 0; i < values.size(); ++i) {
            pout(i) = out[i];
        }
    }

  private:
    typedef enum
    {
        Push = 0,
        Pull,
        PushPull
    } OpType;
    OpType _op;
};
}
}

namespace tensorflow {
REGISTER_KERNEL_BUILDER(Name("PsOp").Device(DEVICE_CPU), ps::tensorflow::PsOp);
REGISTER_KERNEL_BUILDER(Name("PsMyRank").Device(DEVICE_CPU),
                        ps::tensorflow::PsMyRankOp);
}
